package com.wang.wangaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.wang.wangaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @author: Shajia Wang
 * @createTime: 2025/6/4---15:59
 * @description: 处理工具调用的基础代理类，具体实现了 think 和 act 方法，可以用作创建实例的父类
 */

@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {

    /**
     * 可用的工具
     */
    private final ToolCallback[] availableTools;

    /**
     * 保存了工具调用信息的响应
     */
    private ChatResponse toolCallChatResponse;

    /**
     * 工具调用管理者
     */
    private final ToolCallingManager toolCallingManager;

     /**
      * 禁用内置的工具调用机制，自己维护上下文
      */
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        // 禁用 Spring AI 内置的工具调用机制，自己维护选项和消息上下文
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
    }

    /**
     * 处理当前状态并决定下一步行动
     *
     * @return 是否需要执行行动，true表示需要执行，false表示不需要执行
     */
    @Override
    public boolean think() {
        //检验
        if(getNextStepPrompt()!=null && !getNextStepPrompt().isEmpty()){
            UserMessage userMessage = new UserMessage(getNextStepPrompt());
            getMessageList().add(userMessage);
        }
        List<Message> messageList = getMessageList();
        Prompt prompt = new Prompt(messageList, chatOptions);
        try {
            //获取带工具选项的响应
            ChatResponse chatResponse = getChatClient()
                    .prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            this.toolCallChatResponse= chatResponse;
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            //输出提示信息
            String result = assistantMessage.getText();
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            log.info("{} 的思考 {}",  getName(), result);
            log.info("{} 选择了 {} 个工具来使用",getName(),  toolCallList.size());
            String toolCallInfo = toolCallList.stream().
                    map(toolCall -> String.format("工具名称：%s, 参数: %s",
                    toolCall.name(), toolCall.arguments()))
                    .collect(Collectors.joining("\n"));
            log.info(toolCallInfo);
            if (toolCallList.isEmpty()){
                //只有不调用工具时，才记录助手消息
                getMessageList().add(assistantMessage);
                return false;
            }else{
                //需要调用工具时，无需记录助手消息，因为调用工具时会自动记录
                return true;
            }
        }catch (Exception e){
            log.error("{} 思考失败: {}", getName(), e.getMessage());
            getMessageList().add(
                    new AssistantMessage("处理时遇到错误："+e.getMessage()));
            return false;
        }
    }

    /**
     * 执行决定的行动
     *
     * @return 行动执行结果
     */
    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()){
            return "没有需要调用的工具";
        }
        //调用工具
        Prompt prompt = new Prompt(getMessageList(), chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager
                .executeToolCalls(prompt, toolCallChatResponse);
        //记录消息上下文，conversationHistory已经包含了助手消息和工具调用返回的结果
        setMessageList(toolExecutionResult.conversationHistory());
        //当前工具调用的结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String results = toolResponseMessage.
                getResponses().stream()
                .map(response -> "工具" + response.name() + "完成了它的任务！结果：" + response.responseData())
                .collect(Collectors.joining("\n"));
        //判断是否调用了终止工具
        boolean terminateToolCalled = toolResponseMessage.
                getResponses().stream().
                anyMatch(response -> "doTerminate".equals(response.name()));
        if (terminateToolCalled){
            setState(AgentState.FINISHED);
        }

        log.info(results);
        return results;
    }
}

